#include <istream>
#include <vector>
#include <list>
#include <string>
#include <fstream>
#include <iostream>
#include <algorithm>
#include <map>
#include <pthread.h>
#include <boost/format.hpp>
#include <stdlib.h>

#include "components/UPGMA_nearest_neighbor.h"
#include "components/KITSCH_nearest_neighbor.h"
#include "components/bcpp_tree.h"
#include "components/portable_timer.h"
#include "components/perlutil.h"
#include "components/matrix.h"
#include "components/stlhelper.h"
#include "get_arg.h"
#define DEBUG_LEVEL 6
#include "components/debug_func.h"
#include <boost/shared_array.hpp>

using boost::format;
using boost::shared_array;

using std::vector;
using std::ostringstream;

using std::string;
using std::fill;
using std::ofstream;
using std::cout;
using std::cin;
using std::cerr;
using std::ifstream;
using std::map;
using std::pair;
using std::istream;
using std::max;
using std::min;
using std::list;

struct t_ks_pairs
{
	double		ks;
	unsigned gene1;
	unsigned gene2;
	t_ks_pairs(double d, unsigned g1, unsigned g2): ks(d), gene1(g1), gene2(g2){}
};




//________________________________________________________________________________________

//	get_clusters_minimum

//________________________________________________________________________________________
void get_clusters_minimum(	const t_matrix&						matrix,
							const vector<vector<unsigned> >&	cluster_set,
							vector<double>&    					cluster_mins)
{
	debug_func_cerr(5);
	cluster_mins.clear();
	for (unsigned cluster_id = 0; cluster_id < cluster_set.size(); ++cluster_id)
	{
		const vector<unsigned>& curr_cluster = cluster_set[cluster_id];
		double cluster_min = 0.0;

		for (unsigned ii = 0; ii + 1 < curr_cluster.size(); ++ii)
		{
			for (unsigned jj = ii + 1; jj < curr_cluster.size(); ++jj)
			{
				unsigned gene1 = curr_cluster[ii];
				unsigned gene2 = curr_cluster[jj];
				cluster_min = min(matrix[gene1][gene2], cluster_min);
			}
		}
		cluster_mins.push_back(cluster_min);
	}
}

//________________________________________________________________________________________

//	get_symmetrical_matrix_minimum

//________________________________________________________________________________________
void get_symmetrical_matrix_minimum(	const t_matrix&		matrix,
										vector<double>&    	cluster_mins,
										vector<unsigned>&   cluster_min_indices)
{
	debug_func_cerr(5);
	cluster_mins.clear();
	cluster_min_indices.clear();
	for (unsigned ii = 0; ii  < matrix.size(); ++ii)
	{
		double cluster_min = 10000.0;
		unsigned cluster_min_index = 10000;
		for (unsigned jj = 0; jj  < matrix.size(); ++jj)
		{
			if (ii == jj)
				continue;
			if (cluster_min > matrix[ii][jj])
			{
				cluster_min			= matrix[ii][jj];
				cluster_min_index	= jj;
			}
		}
		if (cluster_min_index == 10000)
		{
			cerr << "!!!!!!!!\n\n\n!!!!!!\n!!!!!!\n!!!!!!\n";
		}
		cluster_mins.push_back(cluster_min);
		cluster_min_indices.push_back(cluster_min_index);
	}
}


//________________________________________________________________________________________

//	validate_clusters
//		ignore cluster members whose relationships with other members is entirely inferred

//________________________________________________________________________________________
void validate_clusters(t_matrix&						matrix,
						vector<vector<unsigned> >&		cluster_set,
						const vector<string>&			names)
{
	debug_func_cerr(5);
	vector<vector<unsigned> > new_cluster_set;
	for (unsigned cluster_id = 0; cluster_id < cluster_set.size(); ++cluster_id)
	{
		vector<unsigned>& curr_cluster = cluster_set[cluster_id];
		vector<unsigned>  new_cluster;

		for (unsigned ii = 0; ii + 1 < curr_cluster.size(); ++ii)
		{
			for (unsigned jj = ii + 1; jj < curr_cluster.size(); ++jj)
			{
				unsigned gene1 = curr_cluster[ii];
				unsigned gene2 = curr_cluster[jj];
				if (matrix[gene1][gene2] > 5000.0)
					continue;
				new_cluster.push_back(gene1);
				new_cluster.push_back(gene2);
			}
		}
		// remove duplicated genes
		sort(new_cluster.begin(), new_cluster.end());
		remove_duplicates(new_cluster);

		// do not add to new cluster set if empty
		if (!new_cluster.size())
			continue;
		new_cluster_set.push_back(vector<unsigned>());
		new_cluster_set.back().swap(new_cluster);
	}
	new_cluster_set.swap(cluster_set);
}

//________________________________________________________________________________________

//	partition_into_clusters

//________________________________________________________________________________________
void partition_into_clusters(t_matrix&						matrix,
							double					   		cutoff,
							vector<vector<unsigned> >&		cluster_set)
{
	debug_func_cerr(5);
	//// NB cluster_ids start from 1: zero = unassigned cluster
	vector<unsigned>			gene_to_cluster (matrix.size(), 0);
	vector<list<unsigned> >		gene_clusters(1);
	unsigned					cnt_clusters = 0;

	for (unsigned gene1 = 0; gene1 < matrix.size() - 1 ; ++gene1)
		for (unsigned gene2 = gene1 + 1; gene2 < matrix.size(); ++gene2)
	{
		if (matrix[gene1][gene2] > cutoff)
			continue;
		unsigned& curr_cluster1 = gene_to_cluster[gene1];
		unsigned& curr_cluster2 = gene_to_cluster[gene2];

		// add new cluster if neither is yet clustered
		if (curr_cluster1 + curr_cluster2 == 0)
		{
			// cluster index for these two genes
			curr_cluster1 =
			curr_cluster2 = ++cnt_clusters;

			// these two genes make up a new cluster
			gene_clusters.push_back(list<unsigned>());
			gene_clusters.back().push_back(gene1);
			gene_clusters.back().push_back(gene2);
		}

		// link cluster 2 to cluster 1 if both already have clusters
		else if (curr_cluster1 && curr_cluster2)
		{
			// ignore if both already part of the same cluster
			if (curr_cluster1 == curr_cluster2)
				continue;

			// assign curr_cluster1 to all of curr_cluster2 genes;
			list<unsigned>::iterator genes_iter = gene_clusters[curr_cluster2].begin();
			list<unsigned>::iterator genes_end = gene_clusters[curr_cluster2].end();
			unsigned old_cluster = curr_cluster2;
			while (genes_iter != genes_end)
				gene_to_cluster[*genes_iter++] = curr_cluster1;

			// merge cluster 2 into cluster 1
			gene_clusters[curr_cluster1].splice(gene_clusters[curr_cluster1].end(),
												gene_clusters[old_cluster]);
		}

		// add gene 2 to cluster 1 if gene 2 is unclustered
		else if (curr_cluster1)
			gene_clusters[curr_cluster2 = curr_cluster1].push_back(gene2);

		// add gene 1 to cluster 2 if gene 1 is unclustered
		else
			gene_clusters[curr_cluster1 = curr_cluster2].push_back(gene1);

	}

	//
	// copy non-empty clusters to output set
	//
	for (unsigned i = 0; i < gene_clusters.size(); ++i)
		if (!gene_clusters[i].empty())
			cluster_set.push_back(vector<unsigned>(gene_clusters[i].begin(),
												   gene_clusters[i].end()));
}


//________________________________________________________________________________________

//	scale_data

//________________________________________________________________________________________
void scale_data(t_matrix& matrix, const vector<string>& names,
				double ortholog_ks, double scaling_factor,
				const string& tax_name2)
{
	// constants
	double tax2_ks = (scaling_factor * ortholog_ks) / ( 1.0 + scaling_factor);
	double scaling_factor_inv = 1.0 / scaling_factor;


	// figure out which genes belong to species #2
	vector<double>	is_second_taxa(names.size(), 0.0);
	for (unsigned i = 0; i < names.size(); ++i)
		is_second_taxa.push_back((names[i].find(tax_name2) != string::npos) ? 1.0 : 0.0);

	// scale distance in Ks units from ortholog node to tax2
	for (unsigned i = 0; i < names.size() - 1 ; ++i)
		for (unsigned j = i + 1; j < names.size(); ++j)
		{
			double has_second_taxa_in_pair = is_second_taxa[i] + is_second_taxa[j];

			// scale ks for different orthologs
			double ks = matrix[i][j];
			double scale = std::min(ks, tax2_ks * has_second_taxa_in_pair);
			double delta = ks - scale;
			scale *= scaling_factor_inv;
			matrix[i][j] = matrix[j][i] = scale + delta;
		}
}


//________________________________________________________________________________________

//	read_matrix_from_stream

//________________________________________________________________________________________
void dump_matrix_to_stream(	ostream& outfile,
							t_matrix& matrix,
							vector<string>& names)
{
	outfile << string(15, ' ');
	for (unsigned i = 0; i < matrix.size(); ++i)
		outfile << boost::format ("%|1$-21|") % names[i];
	outfile << "\n";

	for (unsigned i = 0; i < matrix.size(); ++i)
	{
		outfile << boost::format ("%|1$-21|")  % names[i];
		for (unsigned j = 0; j < matrix.size(); ++j)
			outfile << boost::format ("%|1$-21|") %
							matrix[i][j];
		outfile << "\n";
	}
	outfile << "\n";
}
//________________________________________________________________________________________

//	read_matrix_from_stream

//________________________________________________________________________________________
unsigned read_matrix_from_stream(	istream& infile,
									t_matrix& matrix,
									vector<string>& names,
									string& cluster_name)
{


	vector<t_ks_pairs>		ks_pairs;
	map<string, unsigned>	gene_to_index;
	typedef pair<string, unsigned> t_gene_index_pair;
	string line;
	while (getline(infile, line))
	{
		if (!line.length())
		{
			// ignore consecutive blank lines
			if (!names.size())
				continue;
			else if (!ks_pairs.size())
			{
				names.clear();
				gene_to_index.clear();
				continue;
			}

			// blank lines separate clusters
			else
				break;
		}

		if (line[0] == '>')
		{
			cluster_name = line.substr(1);;
			continue;
		}
		vector<string> data;
		perl_split (line, "\t", data, false, 3);
		pair<map<string, unsigned>::iterator, bool> pr;



		//
		// first name in pair
		//
		// map::insert will add item only if it does not already exist
		//     If so, the gene name will correspond to the index of the NEXT
		//     item in "names", i.e. ==names.size()
		unsigned tax1_index = names.size();
		pr = gene_to_index.insert(t_gene_index_pair(data[0],  tax1_index));

		// this name not been encountered before: save name
		if (pr.second)
		{
			names.push_back(data[0]);
		}

		// else use existing value
		else
			tax1_index = (pr.first)->second;


		//
		// second name in pair
		//
		unsigned tax2_index = names.size();
		pr = gene_to_index.insert(t_gene_index_pair(data[1],  tax2_index));

		// this name not been encountered before: save name
		if (pr.second)
		{
			names.push_back(data[1]);
		}
		// else use existing value
		else
			tax2_index = (pr.first)->second;

		double ks = atof(data[2].c_str());
		// ignore those less than 0.0 and greater than 10000
		if (ks >= 0.0 and ks <	 10000.0)
			ks_pairs.push_back(t_ks_pairs(ks, tax1_index, tax2_index));

	}

	// number of species
	unsigned cnt_leaves = names.size();
	if (!cnt_leaves)
		return 0;

	// get data and names
	vector<vector<char> > pair_counts(cnt_leaves);
	for (unsigned i = 0; i < cnt_leaves; ++i)
		pair_counts[i].assign(cnt_leaves, 1);
	unsigned unique_leaves = 0;

	matrix.allocate(cnt_leaves, 10000.0);
	for (unsigned i = 0; i < ks_pairs.size(); ++i)
	{
		unsigned gene1 = ks_pairs[i].gene1;
		unsigned gene2 = ks_pairs[i].gene2;
		unique_leaves += pair_counts[gene1][gene2];
		pair_counts[gene1][gene2] =
		pair_counts[gene2][gene1] = 0;
		matrix[gene1][gene2]=
		matrix[gene2][gene1]= ks_pairs[i].ks;
	}

	return unique_leaves;
}  /* read_matrix_from_stream */



//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

//	print and analyse orthologs

//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
bool open_filename(ofstream& ofs, const string& file_name)
{
	ofs.open(file_name.c_str());
	if (ofs)
		return true;
	cerr << "Error\nCould not open " << file_name << "\n";
	return false;
}
struct t_outputstreams
{

	t_outputstreams(const string& prefix, const std::vector<double>& cutoff) :
			file_prefix(prefix)
	{
		open(cutoff);
		for (unsigned i = 0; i < cutoff.size(); ++i)
			ofs_clusters[i] << "Cluster ID\tCluster size\tNumber of duplicated genes\n";
		ofs_all_nodes	<< "To leaf\tTo root\tGene id\tNode type\n"
						<< string(80, '-') << "\n";

		ofs_full_tree	<< "Branch length\tTo root\tTo leaf\tGene id\tNode type\n"
						<< string(80, '-') << "\n";
	}
	string	 file_prefix;
	ofstream ofs_timing;
	ofstream ofs_leaf_len;
	ofstream ofs_full_tree;
	boost::shared_array <ofstream> ofs_genes;
	ofstream ofs_newhampshire_tree;
	ofstream ofs_all_nodes;
	boost::shared_array <ofstream> ofs_clusters;
	bool open(const std::vector<double>& cutoff);
	bool printout_analysis(
							bcpp_tree::t_tree&						tree,
							const string&								cluster_name,
							const vector<vector<vector<unsigned> > >&	sub_clusters,
							const vector<vector<double> >&    			cluster_min,
							const string&								header,
							const vector<string>&						names );
};





bool t_outputstreams::printout_analysis(
								bcpp_tree::t_tree&						tree,
								const string&								cluster_name,
								const vector<vector<vector<unsigned> > >&	sub_clusters,
								const vector<vector<double> >&    			cluster_min,
								const string&								header,
								const vector<string>&						names )
{
	debug_func_cerr(5);
	// cerr << "dump all nodes\n";
	// dump all nodes
	{
		ofs_all_nodes << "Cluster id = " << cluster_name << "\n";
		bcpp_tree::do_dump_nodes dump_nodes(ofs_all_nodes);
		tree.for_each(dump_nodes);
		ofs_all_nodes << std::endl;
	}

	// cerr << "dump tree diagram\n";
	// print full tree
	{

		ofs_full_tree << "Cluster id = " << cluster_name << "\n";
		bcpp_tree::op_print_tree print_tree(tree, 3, ofs_full_tree);
		print_tree.print();
//		bcpp_tree::do_print_tree_full print_tree(ofs_full_tree);
//		tree.for_each(print_tree);
		ofs_full_tree << std::endl;
	}

	// cerr << "dump nh tree\n";
	// new hampshire tree
	{
		ofs_newhampshire_tree << "Cluster id = " << cluster_name << "\n";
		bcpp_tree::do_newhampshire_format_str nh_str(ofs_newhampshire_tree);
		tree.for_each(nh_str);
		ofs_newhampshire_tree << std::endl;
	}



	// cerr << "dump gene lists\n";
	// print gene sets
	for (unsigned cutoff = 0; cutoff < sub_clusters.size(); ++cutoff)
	{
		unsigned sub_cluster_size = 0;
		{
			for (unsigned i = 0; i < sub_clusters[cutoff].size(); ++i)
			{
				for (unsigned j = 0; j < sub_clusters[cutoff][i].size(); ++j)
				{
					assert(sub_clusters[cutoff][i][j] < names.size());
					ofs_genes[cutoff]
								<< names[sub_clusters[cutoff][i][j]]		<< "\t"
								<< cluster_name					 			<< "\t"
								<< i                						<< std::endl;
				}
				sub_cluster_size += sub_clusters[cutoff][i].size();
			}
		}

		// cerr << "dump cluster stats\n";
		// print cluster stats
		{
			for (unsigned i = 0; i < sub_clusters[cutoff].size(); ++i)
			{
				ofs_clusters[cutoff]<< cluster_name							<< "\t"
									<< sub_cluster_size						<< "\t"
									<< sub_clusters[cutoff][i].size() 		<< std::endl;

			}
		}

	}
	return true;

}



bool t_outputstreams::open(const std::vector<double>& cutoff)
{
	ofs_genes.reset(new ofstream[cutoff.size()]);
	ofs_clusters.reset(new ofstream[cutoff.size()]);
	for (unsigned i = 0; i < cutoff.size(); ++i)
	{
		boost::format fmter ("_%|1$-4.3f|_");
		fmter % cutoff[i];
		string suffix = fmter.str();
		suffix[2]='_';
		if (!open_filename(ofs_genes[i], file_prefix + "genes.list." + suffix))
			return false;
		if (!open_filename(ofs_clusters[i], file_prefix + "clusters.list." + suffix))
			return false;
	}

	return(
			open_filename(ofs_leaf_len, 		file_prefix + "leaf.len"			)	&&
			open_filename(ofs_full_tree, 		file_prefix + "full_topology.tree"	)	&&
			open_filename(ofs_newhampshire_tree,file_prefix + "new_hampshire.tree"	)	&&
			open_filename(ofs_all_nodes, 		file_prefix + "all_nodes.list"		)	&&
			open_filename(ofs_timing, 			file_prefix + "timing"				)

		);
}

struct missing_pairs
{
	unsigned cluster_num;
	unsigned actual_pairs;
	unsigned expected_pairs;
	missing_pairs(unsigned c, unsigned a, unsigned e):
		cluster_num(c), actual_pairs(a), expected_pairs(e){}
};


//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

//	main

//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
int main (int argc, char *argv[])
{
	// we use cout to print progress so flush per call
	cout.setf(std::ios_base::unitbuf);

	std::ios::sync_with_stdio(false);

	t_program_args args;
	if (!process_command_line_options(argc, argv, args))
		return 1;


	//
	// open input file
	//
	istream* datafile = &cin;
	ifstream infile;
	if (args.ks_file.length())
	{
		infile.open(args.ks_file.c_str());
		if (!infile)
		{
			std::cerr << "Error opening [" << args.ks_file << "] \n";
			return 1;
		}
		datafile = &infile;
	}

	//
	//	prepare files for output
	//
	t_outputstreams strms (args.out_file_prefix, args.cutoff);
	strms.ofs_timing <<	"Genes\tSum of squares\tAverage %stddev\tCount Trees\tbest iteration\tTime\n"
					"-----\t--------------\t---------------\t-----------\t--------------\t---------\n";


	unsigned cnt_cluster = 0;
	//
	//	save clusters with missing entries
	//
	vector<missing_pairs> problems;



	{
		t_progress_indicator dots(cerr, 5000);
		while (datafile->good())
		{
			++cnt_cluster;


			//
			//	read data from input file
			//
			t_matrix		matrix;
			vector<string> 	names;
			string			cluster_name;
			unsigned 		cnt_pairs =  read_matrix_from_stream(*datafile, matrix, names,
																	cluster_name);
			if (!cnt_pairs)
				break;

			// save missing pairs
			unsigned expected_pairs = names.size() * (names.size() - 1) / 2;
			if (cnt_pairs != expected_pairs)
				problems.push_back(missing_pairs(cnt_cluster, cnt_pairs, expected_pairs));


			strms.ofs_timing << boost::format
										("%|1$-5|\t")
										% names.size();
			strms.ofs_timing.flush();


			double sum_sqrs				= 0.0;
			double avg_percent_stddev	= 0.0;
			unsigned cnt_trees_examined = 1;
			unsigned best_iteration		= 0;
			t_portable_timer	timer;
			bcpp_tree::t_tree paralog_tree;
			if (args.analysis_type == eUPGMA || args.analysis_type == eUPGMA_KITSCH)
			{
				upgma_tree::upgma_make_tree(paralog_tree, names, matrix,
											args.fill_lacunae, dots);
				// get sum squares for upgma
				bcpp_tree::do_sum_sqrs2 analyse_tree(matrix);
				paralog_tree.for_each(analyse_tree);
				analyse_tree.results(sum_sqrs, avg_percent_stddev);
			}
			else if (args.analysis_type == eKITSCH)
				// kitsch make and evaluate
				kitsch_tree::kitsch_make_tree(paralog_tree, names, matrix,
												args.cnt_random_global_searches,
												args.power,
												best_iteration,
												sum_sqrs, avg_percent_stddev,
												cnt_trees_examined,
												dots);
			// optimize UPGMA tree
			if (args.analysis_type == eUPGMA_KITSCH)
			{
				kitsch_tree::optimize_bcpp_tree(paralog_tree, names, matrix,
													args.cnt_random_global_searches,
													args.power,
													args.keyhole_size,
													best_iteration,
													sum_sqrs,
													avg_percent_stddev,
													cnt_trees_examined,
													dots);
			}
			strms.ofs_timing << boost::format(	"%|1$14|\t"
												"%|2$15|\t"
												"%|3$11|\t"
												"%|4$14|\t"
												"%|5$9|\n")
												% sum_sqrs
												% avg_percent_stddev
												% cnt_trees_examined
												% best_iteration
												% timer.reset();
			strms.ofs_timing.flush();



			// get matrix from constructed tree
			t_matrix matrix_from_tree;
			paralog_tree.to_matrix(matrix_from_tree);

			vector<double>			per_gene_minimum(matrix.size());
			vector<unsigned>		per_gene_minimum_indices(matrix.size());
			get_symmetrical_matrix_minimum(	matrix, per_gene_minimum,
													per_gene_minimum_indices);

			// use minimum values of the original and constructed matrices
			matrix_from_tree.min(matrix);

			// but do not have minima smaller than could be inferred via any relationships
			for (unsigned ii = 0; ii + 1 < matrix_from_tree.size(); ++ii)
			{
				for (unsigned jj = ii + 1; jj < matrix_from_tree.size(); ++jj)
				{
					matrix_from_tree[ii][jj] =
					matrix_from_tree[jj][ii] = max(matrix_from_tree[ii][jj],
										max(per_gene_minimum[ii], per_gene_minimum[jj]));

				}
			}




			vector<vector<vector<unsigned> > >	sub_clusters(args.cutoff.size());
			vector<vector<double> >				cluster_min(args.cutoff.size());
			for (unsigned i = 0; i < sub_clusters.size(); ++i)
			{
				// to partition into sub_clusters
				partition_into_clusters(matrix_from_tree,
										args.cutoff[i],
										sub_clusters[i]);

				// but make sure relationships within sub_clusters are real and not inferred
				// from third parties outside the cluster
				//validate_clusters(matrix, sub_clusters[i], names);
				validate_clusters(matrix_from_tree, sub_clusters[i], names);
				get_clusters_minimum(	matrix_from_tree,
										sub_clusters[i],
										cluster_min[i]);
			}

			// print out less aggressive matrix minimums
			{
				vector<double> final_per_gene_minimum(matrix_from_tree.size());
				vector<unsigned> final_per_gene_min_indices(matrix_from_tree.size());

				get_symmetrical_matrix_minimum(	matrix_from_tree, final_per_gene_minimum,
												final_per_gene_min_indices);
				for (unsigned i = 0; i < final_per_gene_minimum.size(); ++i)
				{
					strms.ofs_leaf_len <<	names[i]					<< "\t" <<
											final_per_gene_minimum[i]	<< "\t" <<
											names[final_per_gene_min_indices[i]] 	<< "\t" <<

											final_per_gene_min_indices[i]
																		<< "\n";
				}
			}


			strms.printout_analysis(paralog_tree, cluster_name, sub_clusters, cluster_min, "",
									names);
		}
	}


	//
	//	print clusters with missing entries
	//
	if (problems.size())
	{
		cerr << "Some matrixes had missing pairs.\n";
		for (unsigned i = 0; i < problems.size(); ++i)
		{
			cout	<< format("\tCluster #%|1$+4|: Only %|2$+4| out of %|3$+4| pairs provided.\n")
					% problems[i].cluster_num
					% problems[i].actual_pairs
					% problems[i].expected_pairs;
		}
	}
	cerr << "\nAll finished!!\n\n";

	return 0;
}

